from contextlib import contextmanager
from urllib.parse import quote_plus

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from core.config import get_config
from sqlalchemy.engine import URL


db_config = get_config('app.datasource')


db_url = URL.create(
    drivername="postgresql+psycopg2",
    username=db_config["username"],
    password=db_config["password"],
    host=db_config["host"],
    port=db_config["port"],
    database=db_config["dbname"]
)
# db_url = f"postgresql+psycopg2://{db_config['username']}:{quote_plus(db_config['password'])}@{db_config['host']}:{db_config['port']}/{db_config['dbname']}"


engine = create_engine(
    db_url,
    pool_size=10,
    max_overflow=20,
    pool_timeout=30,
    pool_recycle=3600,
    echo=False,
    future=True
)

SessionLocal = sessionmaker(bind=engine)

@contextmanager
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# 事务处理
@contextmanager
def get_db_transactional():
    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()